package com.yww.algorithm.binarytree;

import com.yww.datastructure.binarytree.TreeNode;
import com.yww.datastructure.stack.LinkedListStack;
import com.yww.datastructure.stack.Stack;

import java.util.function.Consumer;

/**
 * @author yww
 * @description 二叉树的各种遍历
 * @since 2024/2/8 13:43
 */
public class TreeTraversal {
    public static void main(String[] args) {
        TreeNode<Integer> root = new TreeNode<Integer>(1,
                new TreeNode<>(2,
                        new TreeNode<>(4, null, null), null),
                new TreeNode<>(3,
                        new TreeNode<>(5, null, null),
                        new TreeNode<>(6, null, null)));

        preOrder(root, val -> {
            System.out.print(val + "\t");
        });
        System.out.println();
        inOrder(root, val -> {
            System.out.print(val + "\t");
        });
        System.out.println();
        postOrder(root, val -> {
            System.out.print(val + "\t");
        });
        System.out.println();
        traversal(root, val -> {
            System.out.print(val + "\t");
        }, null, null);
        System.out.println();
        traversal(root, null, val -> {
            System.out.print(val + "\t");
        }, null);
        System.out.println();
        traversal(root, null, null, val -> {
            System.out.print(val + "\t");
        });

    }

    /**
     * 二叉树的前序遍历
     **/
    public static <E> void preOrder(TreeNode<E> cur, Consumer<E> consumer) {
        if (cur == null) {
            return;
        }
        consumer.accept(cur.val); // 先消费
        preOrder(cur.left, consumer); // 递归左
        preOrder(cur.right, consumer); // 递归右
    }

    /**
     * 二叉树的中序遍历
     **/
    public static <E> void inOrder(TreeNode<E> cur, Consumer<E> consumer) {
        if (cur == null) {
            return;
        }
        inOrder(cur.left, consumer); // 递归左
        consumer.accept(cur.val); // 消费
        inOrder(cur.right, consumer); // 递归右
    }

    /**
     * 二叉树的后序遍历
     **/
    public static <E> void postOrder(TreeNode<E> cur, Consumer<E> consumer) {
        if (cur == null) {
            return;
        }
        postOrder(cur.left, consumer); // 递归左
        postOrder(cur.right, consumer); // 递归右
        consumer.accept(cur.val); // 消费
    }

    /**
     * 循环遍历，完成前中后序遍历
     */
    public static <E> void traversal(TreeNode<E> root,
                                     Consumer<E> preOrder,
                                     Consumer<E> inOrder,
                                     Consumer<E> postOrder) {
        TreeNode<E> cur = root;
        Stack<TreeNode<E>> stack = new LinkedListStack<>(); // 寻找返回的路
        TreeNode<E> pop = null; // 当前被弹出元素
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) {
                if (preOrder != null) { // 在处理左节点前，进行前序消费
                    preOrder.accept(cur.val);
                }
                stack.push(cur); // 入栈
                cur = cur.left; // 处理左节点
            } else {
                TreeNode<E> peek = stack.peek(); // 获取即将要弹出的节点
                if (peek.right == null) { //  说明栈顶节点右子树空，无需处理右子树
                    if (inOrder != null) { // 在弹出当前节点前，进行中序消费
                        inOrder.accept(peek.val);
                    }

                    pop = stack.pop(); // 弹出无需处理右子树的节点

                    if (postOrder != null) { // 处理弹出节点，进行后序消费
                        postOrder.accept(pop.val);
                    }
                } else if (peek.right == pop) { // 说明栈顶节点已经处理完成右子树
                    pop = stack.pop(); // 弹出栈顶元素，进行后序消费
                    if (postOrder != null) {
                        postOrder.accept(pop.val);
                    }
                } else { // 说明栈顶节点还右子树要处理
                    if (inOrder != null) { // 在处理右子树前进行中序消费
                        inOrder.accept(peek.val);
                    }
                    cur = peek.right; //处理右子树
                }
            }
        }
    }
}
